import os
import shutil
import subprocess
import tempfile
import json
import traceback
import uuid
import base64
from flask import Flask, request, jsonify, send_file
from datetime import datetime

app = Flask(__name__)

# 修改为实际的工作目录
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
INFERENCE_SCRIPT = os.path.join(BASE_DIR, "indextts")
# 创建用于存储临时输入和输出文件的目录
SAVE_DIR = os.path.join(BASE_DIR, "temp_inputs/")
RESULT_DIR = os.path.join(BASE_DIR, "temp_results/")

# 确保临时目录存在
def ensure_dir_exists(dir_path):
    """确保指定的目录存在，如果不存在则创建"""
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

ensure_dir_exists(SAVE_DIR)
ensure_dir_exists(RESULT_DIR)

def create_unique_dir(base_dir):
    """创建唯一目录并返回路径
    Args:
        base_dir: 基础目录路径
    Returns:
        str: 创建的唯一目录的完整路径
    """
    try:
        os.makedirs(base_dir, exist_ok=True)
        current_dir = datetime.now().strftime("%Y%m%d-%H%M%S%f")
        full_path = os.path.join(base_dir, current_dir)
        os.makedirs(full_path)
        print(f"创建目录: {full_path}")
        return full_path
    except Exception as e:
        print(f"创建目录失败: {base_dir}, 错误: {str(e)}")
        print(traceback.format_exc())
        raise

def generate_random_filename(extension):
    """生成随机文件名，避免文件命名冲突
    Args:
        extension: 文件扩展名，如'.wav'
    Returns:
        str: 生成的随机文件名
    """
    random_name = str(uuid.uuid4())[:8]  # 使用UUID的前8个字符
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    return f"{timestamp}_{random_name}{extension}"

def process_tts(text, voice_path, output_path):
    """处理文本转语音请求，使用IndexTTS模型生成语音
    Args:
        text: 要转换为语音的文本内容
        voice_path: 参考音频文件路径
        output_path: 输出音频文件路径
    Returns:
        str: 生成的音频文件路径
    Raises:
        RuntimeError: 当生成语音失败时抛出异常
    """
    try:
        # 导入IndexTTS进行推理，而不是使用subprocess
        from indextts.infer import IndexTTS
        
        # 初始化IndexTTS模型
        tts = IndexTTS(
            cfg_path=os.path.join(BASE_DIR, "checkpoints/config.yaml"), 
            model_dir=os.path.join(BASE_DIR, "checkpoints"),
            is_fp16=True  # 使用半精度浮点数以提高性能
        )
        
        # 执行推理，将文本转换为语音
        print(f"开始文本转语音处理，文本长度: {len(text)}字符")
        tts.infer(audio_prompt=voice_path, text=text, output_path=output_path, verbose=True)
        
        # 验证输出文件是否存在
        if not os.path.exists(output_path):
            raise RuntimeError(f"生成语音失败，未找到输出文件: {output_path}")
        
        print(f"文本转语音处理成功，输出文件: {output_path}")
        return output_path
        
    except Exception as e:
        print(f"处理文本转语音时出错: {str(e)}")
        print(traceback.format_exc())
        raise

@app.route('/tts', methods=['POST'])
def text_to_speech():
    """处理文本转语音的API端点
    请求参数：
        - audio: 上传的参考音频文件（WAV格式）
        - text: 要转换为语音的文本内容
    返回：
        - status: 处理状态（success或error）
        - audio: 生成的音频数据（base64编码）或错误信息
    """
    # 初始化临时文件路径变量
    temp_voice_path = None  # 临时参考音频文件路径
    temp_output_path = None  # 临时输出音频文件路径
    temp_dir = None  # 临时目录路径
    
    try:
        print("接收到新的文本转语音请求")
        
        # 检查请求数据格式
        if 'audio' not in request.files or 'text' not in request.form:
            return jsonify({'error': '请求中缺少audio文件或text参数'}), 400
        
        # 获取文本内容
        text = request.form['text']
        if not text or len(text.strip()) == 0:
            return jsonify({'error': '文本内容不能为空'}), 400
        
        # 创建临时目录，用于存储上传的参考音频
        temp_dir = create_unique_dir(SAVE_DIR)
        
        # 保存上传的音频文件为随机文件名
        audio_file = request.files['audio']
        audio_extension = '.wav'  # 假设输入音频总是WAV格式
        temp_voice_path = os.path.join(temp_dir, generate_random_filename(audio_extension))
        audio_file.save(temp_voice_path)
        print(f"已保存上传的参考音频: {temp_voice_path}")
        
        # 生成随机的输出文件名
        temp_output_path = os.path.join(RESULT_DIR, generate_random_filename('.wav'))
        
        # 执行文本转语音处理
        output_file = process_tts(text, temp_voice_path, temp_output_path)
        
        # 读取生成的音频文件并返回
        with open(output_file, 'rb') as f:
            audio_data = f.read()
        
        # 将音频数据转换为base64编码，方便在HTTP响应中传输
        audio_base64 = base64.b64encode(audio_data).decode('utf-8')
        
        print(f"文本转语音请求处理成功")
        return jsonify({
            "status": "success",
            "audio": f"data:audio/wav;base64,{audio_base64}"
        })
        
    except Exception as e:
        error_msg = f"处理请求时发生错误: {str(e)}"
        print(error_msg)
        print(traceback.format_exc())
        return jsonify({"error": error_msg}), 500
    finally:
        # 确保在所有情况下都清理临时文件和目录，防止磁盘空间被占用
        # 删除临时参考音频文件
        if temp_voice_path and os.path.exists(temp_voice_path):
            try:
                os.remove(temp_voice_path)
                print(f"已删除临时参考音频文件: {temp_voice_path}")
            except Exception as e:
                print(f"删除临时参考音频文件失败: {temp_voice_path}, 错误: {str(e)}")
                
        # 删除临时输出音频文件
        if temp_output_path and os.path.exists(temp_output_path):
            try:
                os.remove(temp_output_path)
                print(f"已删除临时输出音频文件: {temp_output_path}")
            except Exception as e:
                print(f"删除临时输出音频文件失败: {temp_output_path}, 错误: {str(e)}")
                
        # 删除临时目录
        if temp_dir and os.path.exists(temp_dir) and os.path.isdir(temp_dir):
            try:
                shutil.rmtree(temp_dir)
                print(f"已删除临时目录: {temp_dir}")
            except Exception as e:
                print(f"删除临时目录失败: {temp_dir}, 错误: {str(e)}")

if __name__ == '__main__':
    print(f"IndexTTS 服务器启动，监听端口: 6566")
    print(f"基础目录: {BASE_DIR}")
    print(f"输入临时目录: {SAVE_DIR}")
    print(f"输出临时目录: {RESULT_DIR}")
    # 启动Flask服务器，监听所有网络接口的6566端口
    app.run(host='0.0.0.0', port=6566)